#include "menu.hpp"
#include "svnrev.h"
#include "loader/sys.h"
#include "http.h"
#include "pngu.h"

//#include "loader/fs.h"

#include "unzip/unzip.h"
#include "unzip/miniunz.h"

//#include <network.h>
#include "networkops.h"

#include "gecko.h"
#include <fstream>
#include "lockMutex.hpp"

#include <ogc/lwp_watchdog.h>
#include <time.h>

#define UPDATE_URL_VERSION	"http://update.wiiflow.org/txt/versions.txt"

using namespace std;

static block download = { 0, 0 };

void CMenu::_hideDownload(bool instant)
{
	m_btnMgr.hide(m_downloadLblTitle, instant);
	m_btnMgr.hide(m_downloadBtnCancel, instant);
	m_btnMgr.hide(m_downloadPBar, instant);
	m_btnMgr.hide(m_downloadLblMessage[0], 0, 0, -2.f, 0.f, instant);
	m_btnMgr.hide(m_downloadLblMessage[1], 0, 0, -2.f, 0.f, instant);
}

void CMenu::_showDownload(void)
{
	m_btnMgr.show(m_downloadLblTitle);
	m_btnMgr.show(m_downloadBtnCancel);
}

void CMenu::_setThrdMsg(const wstringEx &msg, float progress)
{
	if (m_thrdStop) return;
	if (msg != L"...") m_thrdMessage = msg;
	m_thrdMessageAdded = true;
	m_thrdProgress = progress;
}

bool CMenu::_downloadProgress(void *obj, int size, int position)
{
	CMenu *m = (CMenu *)obj;
	LWP_MutexLock(m->m_mutex);
	m->_setThrdMsg(L"...", m->m_thrdStep + m->m_thrdStepLen * ((float)position / (float)size));
	LWP_MutexUnlock(m->m_mutex);
	return !m->m_thrdStop;
}

void CMenu::_download()
{
	lwp_t thread = LWP_THREAD_NULL;
	int msg = 0;
	wstringEx prevMsg;
	
	SetupInput();
	_showDownload();
	m_btnMgr.setText(m_downloadBtnCancel, _t("dl1", L"Cancel"));
	m_thrdStop = false;
	m_thrdMessageAdded = false;

	while (true)
	{
		_mainLoopCommon(m_thrdWorking);
		if ((BTN_HOME_PRESSED || BTN_B_PRESSED || m_exit) && !m_thrdWorking)
			break;
		else if (BTN_UP_PRESSED)
			m_btnMgr.up();
		else if (BTN_DOWN_PRESSED)
			m_btnMgr.down();
		if (BTN_A_PRESSED && !(m_thrdWorking && m_thrdStop))
		{
			if (m_btnMgr.selected(m_downloadBtnCancel))
			{
				LockMutex lock(m_mutex);
				m_thrdStop = true;
				m_thrdMessageAdded = true;
				m_thrdMessage = _t("dlmsg6", L"Canceling...");
			}
		}
		if (Sys_Exiting())
		{
			LockMutex lock(m_mutex);
			m_thrdStop = true;
			m_thrdMessageAdded = true;
			m_thrdMessage = _t("dlmsg6", L"Canceling...");
			m_thrdWorking = false;
		}
		// 
		if (m_thrdMessageAdded)
		{
			LockMutex lock(m_mutex);
			m_thrdMessageAdded = false;
			m_btnMgr.setProgress(m_downloadPBar, m_thrdProgress);
			if (m_thrdProgress == 1.f) m_btnMgr.setText(m_downloadBtnCancel, _t("dl2", L"Back"));

			if (prevMsg != m_thrdMessage)
			{
				prevMsg = m_thrdMessage;
				m_btnMgr.setText(m_downloadLblMessage[msg], m_thrdMessage, false);
				m_btnMgr.hide(m_downloadLblMessage[msg], 0, 0, -1.f, -1.f, true);
				m_btnMgr.show(m_downloadLblMessage[msg]);
				msg ^= 1;
				m_btnMgr.hide(m_downloadLblMessage[msg], 0, 0, -1.f, -1.f);
			}
		}
		if (m_thrdStop && !m_thrdWorking)
			break;
	}
	if (thread != LWP_THREAD_NULL)
	{
		LWP_JoinThread(thread, NULL);
		thread = LWP_THREAD_NULL;
	}
	_hideDownload();
}

void CMenu::_initDownloadMenu(CMenu::SThemeData &theme)
{
	m_downloadLblTitle = _addLabel(theme, "DOWNLOAD/TITLE", theme.titleFont, L"", 20, 30, 600, 60, theme.titleFontColor, FTGX_JUSTIFY_CENTER | FTGX_ALIGN_MIDDLE);
	m_downloadPBar = _addProgressBar(theme, "DOWNLOAD/PROGRESS_BAR", 40, 200, 560, 20);
	m_downloadBtnCancel = _addButton(theme, "DOWNLOAD/CANCEL_BTN", theme.btnFont, L"", 420, 410, 200, 56, theme.btnFontColor);
	m_downloadLblMessage[0] = _addLabel(theme, "DOWNLOAD/MESSAGE1", theme.lblFont, L"", 40, 228, 560, 100, theme.txtFontColor, FTGX_JUSTIFY_LEFT | FTGX_ALIGN_TOP);
	m_downloadLblMessage[1] = _addLabel(theme, "DOWNLOAD/MESSAGE2", theme.lblFont, L"", 40, 228, 560, 100, theme.txtFontColor, FTGX_JUSTIFY_LEFT | FTGX_ALIGN_TOP);
	// 
	_setHideAnim(m_downloadLblTitle, "DOWNLOAD/TITLE", 0, 0, -2.f, 0.f);
	_setHideAnim(m_downloadPBar, "DOWNLOAD/PROGRESS_BAR", 0, 0, -2.f, 0.f);
	_setHideAnim(m_downloadBtnCancel, "DOWNLOAD/CANCEL_BTN", 0, 0, -2.f, 0.f);
	_hideDownload(true);
	_textDownload();
}

void CMenu::_textDownload(void)
{
	m_btnMgr.setText(m_downloadBtnCancel, _t("dl1", L"Cancel"));
	m_btnMgr.setText(m_downloadLblTitle, _t("dl5", L"Download"));
}

s8 CMenu::_versionTxtDownloaderInit(CMenu *m) //Handler to download versions txt file
{
	if (!m->m_thrdWorking) return 0;
	return m->_versionTxtDownloader();
}

s8 CMenu::_versionTxtDownloader() // code to download new version txt file
{
/* 	u32 bufferSize = 0x001000;	// Maximum download size 4kb
	SmartBuf buffer = smartCoverAlloc(bufferSize);
	if (!buffer)
	{
		LWP_MutexLock(m_mutex);
		_setThrdMsg(L"Not enough memory", 1.f);
		LWP_MutexUnlock(m_mutex);
		m_thrdWorking = false;
		return 0;
	} */

	LWP_MutexLock(m_mutex);
	_setThrdMsg(_t("dlmsg1", L"Initializing network..."), 0.f);
	LWP_MutexUnlock(m_mutex);
	
	if (!IsNetworkInit())
	{
		LWP_MutexLock(m_mutex);
		_setThrdMsg(_t("dlmsg2", L"Network initialization failed!"), 1.f);
		LWP_MutexUnlock(m_mutex);
	}
	else
	{
		// DLoad txt file
		LWP_MutexLock(m_mutex);
		_setThrdMsg(_t("dlmsg11", L"Downloading..."), 0.2f);
		LWP_MutexUnlock(m_mutex);

		m_thrdStep = 0.2f;
		m_thrdStepLen = 0.9f - 0.2f;
		gprintf("TXT update URL: %s\n\n", m_cfg.getString("GENERAL", "updatetxturl", UPDATE_URL_VERSION).c_str());
		download = downloadfile(m_cfg.getString("GENERAL", "updatetxturl", UPDATE_URL_VERSION).c_str());
		if (download.data == 0 || download.size < 19)
		{
			LWP_MutexLock(m_mutex);
			_setThrdMsg(_t("dlmsg20", L"No version information found."), 1.f); // TODO: Check for 16
			LWP_MutexUnlock(m_mutex);
		}
		else
		{
			// txt download finished, now save file
			LWP_MutexLock(m_mutex);
			_setThrdMsg(_t("dlmsg13", L"Saving..."), 0.9f);
			LWP_MutexUnlock(m_mutex);			
			
			FILE *file = fopen(m_ver.c_str(), "wb");
			if (file != NULL)
			{
				fwrite(download.data, 1, download.size, file);
				SAFE_CLOSE(file);

				// version file valid, check for version with SVN_REV
				int svnrev = atoi(SVN_REV);
				gprintf("Installed Version: %d\n", svnrev);
				m_version.load(m_ver.c_str());
				int rev = m_version.getInt("GENERAL", "version", 0);
				gprintf("Latest available Version: %d\n", rev);
				if (svnrev < rev)
				{
					// new version available
					LWP_MutexLock(m_mutex);
					_setThrdMsg(_t("dlmsg19", L"New update available!"), 1.f);
					LWP_MutexUnlock(m_mutex);
				}
				else
				{
					// no new version available
					LWP_MutexLock(m_mutex);
					_setThrdMsg(_t("dlmsg17", L"No new updates found."), 1.f);
					LWP_MutexUnlock(m_mutex);
				}
			}
			else
			{
				LWP_MutexLock(m_mutex);
				_setThrdMsg(_t("dlmsg15", L"Saving failed!"), 1.f);
				LWP_MutexUnlock(m_mutex);
			}

		}
	}
	m_thrdWorking = false;
	return 0;
}

s8 CMenu::_versionDownloaderInit(CMenu *m) //Handler to download new dol
{
	if (!m->m_thrdWorking) return 0;
	return m->_versionDownloader();
}

s8 CMenu::_versionDownloader() // code to download new version
{
	char dol_backup[33];
	strcpy(dol_backup, m_dol.c_str());
	strcat(dol_backup, ".backup");

/* 	u32 updateZipSize = m_data_update_size > m_app_update_size ? m_data_update_size : m_app_update_size;
	u32 sizeToBuffer = updateZipSize > m_dol_update_size ? updateZipSize : m_dol_update_size;
	//sizeToBuffer is in bits, if it is not set in the ini, alloc 4MB
	sizeToBuffer +=2000; //Add header room. */
	if (m_dol_update_size == 0) m_dol_update_size = 1*0x400000;
	if (m_app_update_size == 0) m_app_update_size = 1*0x400000;
	if (m_data_update_size == 0) m_data_update_size = 1*0x400000;
/* 	 if(sizeToBuffer == 2000)  sizeToBuffer = 1*0x400000; */	// Buffer for size of the biggest file.

	
	// check for existing dol
    ifstream filestr;
	gprintf("DOL Path: %s\n", m_dol.c_str());
    filestr.open(m_dol.c_str());
    if (filestr.fail())
	{
		filestr.close();
		rename(dol_backup, m_dol.c_str());
		filestr.open(m_dol.c_str());
		if (filestr.fail())
		{
			LWP_MutexLock(m_mutex);
			_setThrdMsg(_t("dlmsg18", L"boot.dol not found at default path!"), 1.f);
			LWP_MutexUnlock(m_mutex);
		}
		filestr.close();
		sleep(3);
		m_thrdWorking = false;
		return 0;
	}
    filestr.close();

/* 	u32 bufferSize = 0x400000;	// Buffer for size of the biggest file.
	SmartBuf buffer = smartAnyAlloc(bufferSize);
	if (!buffer)
	{
		LWP_MutexLock(m_mutex);
		_setThrdMsg(L"Not enough memory!", 1.f);
		LWP_MutexUnlock(m_mutex);
		sleep(3);
		m_thrdWorking = false;
		return 0;
	} */

	LWP_MutexLock(m_mutex);
	_setThrdMsg(_t("dlmsg1", L"Initializing network..."), 0.f);
	LWP_MutexUnlock(m_mutex);
	
	if (!IsNetworkInit())
	{
		LWP_MutexLock(m_mutex);
		_setThrdMsg(_t("dlmsg2", L"Network initialization failed!"), 1.f);
		LWP_MutexUnlock(m_mutex);
		sleep(3);
		m_thrdWorking = false;
		return 0;
	}
	else
	{
		// Load actual file
		LWP_MutexLock(m_mutex);
		_setThrdMsg(_t("dlmsg22", L"Updating application directory..."), 0.2f);
		LWP_MutexUnlock(m_mutex);

		m_thrdStep = 0.2f;
		m_thrdStepLen = 0.9f - 0.2f;
		gprintf("App Update URL: %s\n", m_app_update_url);
		gprintf("Data Update URL: %s\n", m_app_update_url);
		
		bool zipfile = true;
		if(!m_update_dolOnly)
			download = downloadfile(m_app_update_url);
		if (m_update_dolOnly || download.data == 0 || download.size < m_app_update_size)
		{
 			// Replace zip with dol
			strcpy(strrchr(m_app_update_url, '.'), ".dol"); 
			zipfile = false;
	
			download = downloadfile(m_app_update_url);
			if (download.data == 0 || download.size < m_dol_update_size)
			{
				download = downloadfile(m_old_update_url);
				if (download.data == 0 || download.size < m_dol_update_size)
				{
					LWP_MutexLock(m_mutex);
					_setThrdMsg(_t("dlmsg12", L"Download failed!"), 1.f);
					LWP_MutexUnlock(m_mutex);
					sleep(3);
					m_thrdWorking = false;
					return 0;
				}
			}
		}
		
		if (download.data > 0)
		{
			// download finished, backup boot.dol and write new files.
			LWP_MutexLock(m_mutex);
			_setThrdMsg(_t("dlmsg13", L"Saving..."), 0.8f);
			LWP_MutexUnlock(m_mutex);			
			
			remove(dol_backup);
			rename(m_dol.c_str(), dol_backup);

			FILE *file =  NULL;
			if (zipfile)
			{
				remove(m_app_update_zip.c_str());

				file = fopen(m_app_update_zip.c_str(), "wb");
				if (file != NULL)
				{
					fwrite(download.data, 1, download.size, file);
					SAFE_CLOSE(file);
					
					LWP_MutexLock(m_mutex);
					_setThrdMsg(_t("dlmsg24", L"Extracting..."), 0.8f);
					LWP_MutexUnlock(m_mutex);
					
					unzFile unzfile = unzOpen(m_app_update_zip.c_str());
					if (unzfile == NULL) goto fail;

					int ret = extractZip(unzfile, 0, 1, NULL, m_appDir.c_str());
					unzClose(unzfile);
					if (ret == UNZ_OK)
					{
						//Update apps dir succeeded, try to update the data dir.
						download.data = NULL;
						download.size = 0;

						LWP_MutexLock(m_mutex);
						_setThrdMsg(_t("dlmsg23", L"Updating data directory..."), 0.2f);
						LWP_MutexUnlock(m_mutex);

						download = downloadfile(m_data_update_url);
						if (download.data == 0 || download.size < m_data_update_size)
						{
							LWP_MutexLock(m_mutex);
							_setThrdMsg(_t("dlmsg12", L"Download failed!"), 1.f);
							LWP_MutexUnlock(m_mutex);
							goto success;
						}
						
						if (download.data > 0)
						{
							// download finished, write new files.
							LWP_MutexLock(m_mutex);
							_setThrdMsg(_t("dlmsg13", L"Saving..."), 0.9f);
							LWP_MutexUnlock(m_mutex);
							
							remove(m_data_update_zip.c_str());

							file = fopen(m_data_update_zip.c_str(), "wb");
							if (file != NULL)
							{
								fwrite(download.data, 1, download.size, file);
								SAFE_CLOSE(file);
								
								LWP_MutexLock(m_mutex);
								_setThrdMsg(_t("dlmsg24", L"Extracting..."), 0.8f);
								LWP_MutexUnlock(m_mutex);
								
								unzFile unzfile = unzOpen(m_data_update_zip.c_str());
								if (unzfile == NULL)
									goto success;

								int ret = extractZip(unzfile, 0, 1, NULL, m_dataDir.c_str());
								unzClose(unzfile);

								if (ret != UNZ_OK)
								{
									LWP_MutexLock(m_mutex);
									_setThrdMsg(_t("dlmsg15", L"Saving failed!"), 1.f);
									LWP_MutexUnlock(m_mutex);
								}
							}
						}
						goto success;
					}
					else
						goto fail;
				}
				else
					goto fail;
			}
			else
			{
				file = fopen(m_dol.c_str(), "wb");
				if (file == NULL)
					goto fail;

				fwrite(download.data, 1, download.size, file);
				SAFE_CLOSE(file);
				goto success;
			}
		}
	}
success:
	LWP_MutexLock(m_mutex);
	_setThrdMsg(_t("dlmsg21", L"WiiFlow will now exit to allow the update to take effect."), 1.f);
	LWP_MutexUnlock(m_mutex);

    filestr.open(m_dol.c_str());
    if (filestr.fail())
	{
		LWP_MutexLock(m_mutex);
		_setThrdMsg(_t("dlmsg25", L"Extraction must have failed! Renaming the backup to boot.dol"), 1.f);
		LWP_MutexUnlock(m_mutex);
		rename(dol_backup, m_dol.c_str());
	}
    filestr.close();

	m_exit = true;
	goto out;

fail:
	rename(dol_backup, m_dol.c_str());
	LWP_MutexLock(m_mutex);
	_setThrdMsg(_t("dlmsg15", L"Saving failed!"), 1.f);
	LWP_MutexUnlock(m_mutex);
out:
	sleep(3);
	m_thrdWorking = false;
	return 0;
}